
# coding: utf-8

import torch


im=5
k=3
s=1
out_width=im-k+1
img = torch.ones(5,5)
kernel = torch.randn(3,3)
zeros= torch.zeros(5,5)
KERNEL = torch.zeros(out_width**2,im**2)
PIXEL = torch.zeros(out_width**2,k**2)
print(zeros.shape)
for i in range(0,out_width):
    for j in range(0,out_width):
        zeros[i:k+i,j:k+j]=kernel
        #print(zeros.view(25))
        KERNEL[i*out_width+j,:] = zeros.view(25)
        zeros= torch.zeros(5,5)
        #print(img[i:k+i,j:k+j].reshape(9))
        PIXEL[i*out_width+j,:] =img[i:k+i,j:k+j].reshape(9) 
print(KERNEL.shape)
print(PIXEL.shape)


# In[41]:


print(KERNEL.t().shape)


# In[44]:


img.reshape(1,25).mm(KERNEL.t())


# In[45]:


kernel.reshape(1,9).mm(PIXEL.t())


# # Z = sigmoid(O)
# # O = conv(img,kernel)
# # loss = 0.5*(Z-Y)^2
#Y=torch.nn.functional.one_hot(torch.tensor(2),9)
Y=torch.ones(1,9)*0.8
for iter in range(0,3000):
    # forward
    O=img.reshape(1,25).mm(KERNEL.t())
    Z = torch.sigmoid(O)
    
    E = Z-Y
    if iter%99==0:
        print("iter=%s,loss = %s"%(iter,((E**2).sum()).numpy()))
        print(Z,Y)
    
    
    # # DLoss/Dimg = DLoss/DZ x DZ/DO x DO/Dimg
    ((E**(1-Z)).mm(KERNEL)).shape
    # # DLoss/Dkernel = DLoss/DZ x DZ/DO x DO/Dkernel
    Dkernel=E*Z*(1-Z).mm(PIXEL)
    kernel = kernel-0.001* Dkernel.view(3,3)
    for i in range(0,out_width):
        for j in range(0,out_width):
            zeros[i:k+i,j:k+j]=kernel
            #print(zeros.view(25))
            KERNEL[i*out_width+j,:] = zeros.view(25)
            zeros= torch.zeros(5,5)
            #print(img[i:k+i,j:k+j].reshape(9))
            PIXEL[i*out_width+j,:] =img[i:k+i,j:k+j].reshape(9) 
    
